import numpy as np
import matplotlib.pyplot as plt


def nms(bboxes, confidence_score, iou_threshold):
    bboxes = np.array(bboxes)
    scores = np.array(confidence_score)
    # 按置信度倒序
    ordered = scores.argsort()

    x_min = bboxes[:, 0]
    y_min = bboxes[:, 1]
    x_max = bboxes[:, 2]
    y_max = bboxes[:, 3]
    # 每个检测框的面积
    area = (x_max - x_min) * (y_max - y_min)

    # 挑选出的检测框（保留框）
    picked_bbox = []
    picked_score = []

    while ordered.size > 0:
        index = ordered[-1]
        picked_bbox.append(bboxes[index])
        picked_score.append(scores[index])
        ix_min = np.maximum(x_min[index], x_min[ordered[:-1]])
        iy_min = np.maximum(y_min[index], y_min[ordered[:-1]])
        ix_max = np.minimum(x_max[index], x_max[ordered[:-1]])
        iy_max = np.minimum(y_max[index], y_max[ordered[:-1]])
        i_area = np.maximum(0, ix_max - ix_min) * np.maximum(0, iy_max - iy_min)
        iou = i_area / (area[index] + area[ordered[:-1]] - i_area)
        # 移除IOU大于阈值的检测框
        ordered = ordered[:-1][iou <= iou_threshold]
    return picked_bbox, picked_score


def do():
    show_nms(nms, 0.3)


def show_nms(nms_fn, iou_threshold):
    # 固定测试数据（4个分类，每个分类包含重叠/非重叠检测框）
    test_data = [
        # 类别0：3个框（两个高度重叠，一个独立）
        (0, [100, 100, 200, 200], 0.9),  # 高置信度框
        (0, [120, 120, 220, 220], 0.75),  # 低置信度框（应被抑制）
        (0, [300, 100, 400, 200], 0.85),  # 独立框（保留）

        # 类别1：4个框（三个部分重叠，一个远离）
        (1, [150, 300, 250, 400], 0.88),  # 高置信度框
        (1, [160, 310, 260, 410], 0.82),  # 中置信度框（应被抑制）
        (1, [170, 320, 270, 420], 0.78),  # 低置信度框（应被抑制）
        (1, [400, 300, 500, 400], 0.8),  # 独立框（保留）

        # 类别2：3个框（两个部分重叠，一个远离）
        (2, [100, 400, 200, 500], 0.92),  # 高置信度框
        (2, [120, 420, 220, 520], 0.86),  # 中置信度框（应被抑制）
        (2, [300, 400, 400, 500], 0.89),  # 独立框（保留）

        # 类别3：5个框（多个重叠+独立）
        (3, [200, 100, 300, 200], 0.95),  # 高置信度框
        (3, [220, 120, 320, 220], 0.89),  # 中置信度框（应被抑制）
        (3, [240, 140, 340, 240], 0.83),  # 低置信度框（应被抑制）
        (3, [400, 100, 500, 200], 0.91),  # 独立框（保留）
        (3, [420, 120, 520, 220], 0.87),  # 独立框（保留）
    ]

    # 拆分数据
    all_bboxes = [box[1] for box in test_data]
    all_scores = [box[2] for box in test_data]
    all_classes = [box[0] for box in test_data]

    # 颜色列表（对应4个类别）
    colors = ['r', 'g', 'b', 'purple']

    # 应用NMS
    filtered_bboxes, filtered_scores = nms_fn(all_bboxes, all_scores, iou_threshold)

    # 找出NMS后的类别对应关系
    filtered_indices = []
    for bbox in filtered_bboxes:
        for i, b in enumerate(all_bboxes):
            if np.array_equal(bbox, b):
                filtered_indices.append(i)
                break
    filtered_classes = [test_data[i][0] for i in filtered_indices]

    # 可视化
    plt.figure(figsize=(12, 6))

    # 原始检测框
    plt.subplot(1, 2, 1)
    plt.imshow(np.zeros((512, 512, 3)), cmap='gray')
    plt.title("Original Detections")

    for i, (bbox, cls_id) in enumerate(zip(all_bboxes, all_classes)):
        x1, y1, x2, y2 = bbox
        rect = plt.Rectangle((x1, y1), x2 - x1, y2 - y1,
                             linewidth=2, edgecolor=colors[cls_id], facecolor='none')
        plt.gca().add_patch(rect)

    # NMS后检测框
    plt.subplot(1, 2, 2)
    plt.imshow(np.zeros((512, 512, 3)), cmap='gray')
    plt.title("After NMS")

    for i, (bbox, cls_id) in enumerate(zip(filtered_bboxes, filtered_classes)):
        x1, y1, x2, y2 = bbox
        rect = plt.Rectangle((x1, y1), x2 - x1, y2 - y1,
                             linewidth=2, edgecolor=colors[cls_id], facecolor='none')
        plt.gca().add_patch(rect)

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    do()
